package DianShang_2024.ds_03.feature

import org.apache.spark.ml.feature.{OneHotEncoder, OneHotEncoderModel, StringIndexer, StringIndexerModel}
import org.apache.spark.ml.functions.vector_to_array
import org.apache.spark.sql.functions.col
import org.apache.spark.sql.{DataFrame, SparkSession}

object test02 {
  def main(args: Array[String]): Unit = {
    //  1.创建 SparkSession 对象
    val spark: SparkSession = SparkSession
      .builder()
      .master("local[*]")
      .appName("One_Hot")
      .getOrCreate()

    //  2.创建一个示例数据集，格式为 DataFrame
    import spark.implicits._

    val data: DataFrame = Seq(
      (0, "red"),
      (1, "blue"),
      (2, "green"),
      (3, "red"),
      (4, "red"),
      (5, "blue")
    ).toDF("id", "color")

    //  3.创建 StringIndexer 对象，对特征进行索引编码
    val indexer: StringIndexerModel = new StringIndexer()
      .setInputCol("color")
      .setOutputCol("index_color")
      .fit(data)

    val indexData: DataFrame = indexer.transform(data)

    //  4.使用 OneHotEncoder 对特征进行 One-Hot 编码
    val encoder: OneHotEncoder = new OneHotEncoder() // OneHotEncoderEstimator()
      .setInputCol("index_color")
      .setOutputCol("encoded_color")
      .setDropLast(false)

    val encoderModel: OneHotEncoderModel = encoder.fit(indexData)

    val encoderData: DataFrame = encoderModel.transform(indexData)

    //  5.输出结果
    encoderData.show(false)

    //    6.进行向量拆分
    encoderData.withColumn("new_encoded_color", vector_to_array(col("encoded_color")))
      .show(false)
  }

}
